import win32gui
from ctypes import windll, byref, c_ubyte
from ctypes.wintypes import RECT
import numpy as np
import cv2

# 排除缩放干扰
windll.user32.SetProcessDPIAware()

class Windowcapture:
    def __init__(self, window_title):
        """
        初始化窗口捕获类。
        
        :param window_title: 窗口标题（如 "雷电模拟器"）
        """
        self.handle = win32gui.FindWindow(None, window_title)
        if self.handle == 0:
            raise ValueError(f"窗口 '{window_title}' 未找到!")

    def capture(self):
        """
        捕获窗口的当前画面。
        
        :return: 返回一个 numpy.ndarray 格式的图像（BGRA 格式）
        """
        # 获取窗口大小
        r = RECT()
        windll.user32.GetWindowRect(self.handle, byref(r))
        width = r.right - r.left
        height = r.bottom - r.top

        # 开始截图
        dc = windll.user32.GetDC(self.handle)
        cdc = windll.gdi32.CreateCompatibleDC(dc)
        bitmap = windll.gdi32.CreateCompatibleBitmap(dc, width, height)
        windll.gdi32.SelectObject(cdc, bitmap)
        windll.gdi32.BitBlt(cdc, 0, 0, width, height, dc, 0, 0, 0x00CC0020)  # SRCCOPY

        # 获取截图数据
        total_bytes = width * height * 4
        buffer = bytearray(total_bytes)
        byte_array = (c_ubyte * total_bytes).from_buffer(buffer)
        windll.gdi32.GetBitmapBits(bitmap, total_bytes, byte_array)

        # 释放资源
        windll.gdi32.DeleteObject(bitmap)
        windll.gdi32.DeleteObject(cdc)
        windll.user32.ReleaseDC(self.handle, dc)

        # 返回 numpy.ndarray 格式的图像
        return np.frombuffer(buffer, dtype=np.uint8).reshape(height, width, 4)